package com.hwadee.backend.util;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.hwadee.backend.config.DoubaoConfig;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;

@Component
public class DoubaoApiUtil {
    @Autowired
    private DoubaoConfig doubaoConfig;
    private static final String API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions";

    public static class ChatRequestBuilder {
        private String model;
        private JSONArray messages = new JSONArray();
        private Float temperature = 1f;
        private Float topP = 0.7f;
        private Integer maxTokens;
        private Object stop;

        public ChatRequestBuilder setModel(String model) {
            this.model = model;
            return this;
        }

        public ChatRequestBuilder addMessage(String role, String content) {
            JSONObject message = new JSONObject();
            message.put("role", role);
            message.put("content", content);
            this.messages.add(message);
            return this;
        }

        public ChatRequestBuilder setTemperature(Float temperature) {
            this.temperature = temperature;
            return this;
        }

        public ChatRequestBuilder setTopP(Float topP) {
            this.topP = topP;
            return this;
        }

        public ChatRequestBuilder setMaxTokens(Integer maxTokens) {
            this.maxTokens = maxTokens;
            return this;
        }

        public ChatRequestBuilder setStop(Object stop) {
            this.stop = stop;
            return this;
        }

        public JSONObject build() {
            JSONObject requestBody = new JSONObject();
            requestBody.put("model", model);
            requestBody.put("messages", messages);
            requestBody.put("temperature", temperature);
            requestBody.put("top_p", topP);
            if (maxTokens != null) {
                requestBody.put("max_tokens", maxTokens);
            }
            if (stop != null) {
                requestBody.put("stop", stop);
            }
            return requestBody;
        }
    }

    public String chatWithDoubao(ChatRequestBuilder builder) throws IOException {
        CloseableHttpClient httpClient = HttpClients.createDefault();
        HttpPost httpPost = new HttpPost(API_URL);

        // 设置请求头
        httpPost.setHeader("Content-Type", "application/json");
        httpPost.setHeader("Authorization", "Bearer " + doubaoConfig.getApiKey());

        // 构建请求体
        JSONObject requestBody = builder.build();
        httpPost.setEntity(new StringEntity(requestBody.toString(), "UTF-8"));

        // 发送请求并获取响应
        try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
            int statusCode = response.getStatusLine().getStatusCode();
            if (statusCode >= 200 && statusCode < 300) {
                HttpEntity entity = response.getEntity();
                if (entity != null) {
                    String result = EntityUtils.toString(entity);
                    JSONObject responseJson = JSONObject.parseObject(result);
                    JSONArray choices = responseJson.getJSONArray("choices");
                    if (choices != null && !choices.isEmpty()) {
                        JSONObject firstChoice = choices.getJSONObject(0);
                        JSONObject messageObj = firstChoice.getJSONObject("message");
                        return messageObj.getString("content");
                    }
                }
            } else {
                HttpEntity entity = response.getEntity();
                String errorResponse = entity != null ? EntityUtils.toString(entity) : "无错误信息";
                throw new IOException("请求失败，状态码: " + statusCode + "，响应: " + errorResponse);
            }
        }
        return "请求失败，请检查参数或网络连接。";
    }
}